from Bio import SeqIO
from Bio import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Align.Applications import ClustalwCommandline
from Bio import AlignIO
from Bio.SeqUtils import GC
import sys, os

def file_extension(filename):
  file_ext = "not_valid"
  if filename.endswith("fasta") or filename.endswith("fa"):
    file_ext = "fasta"
  elif filename.endswith("gbk") or filename.endswith("gb"): 
    file_ext = "genbank"
  return file_ext

def read_seqs(filenames): 
  seq_records = []
  for filename in filenames:
    file_ext = file_extension(filename) 
    if file_ext != "not_valid":
      for seq_record in SeqIO.parse(filename, file_ext):
        seq_records.append(seq_record)
    else:
      sys.stderr.write("Skip file " + filename + " as it have not valid format")
  return (seq_records) 

def get_gc_content(seq_record):
  return str(GC(seq_record.seq))

def transcription(seq_record):
  return str(Seq.transcribe(seq_record.seq))

def translation(seq_record):
  return str(Seq.translate(seq_record.seq))

def multiple_align(seq_records, phylip_filename):
    if len(seq_records) == 1:
      raise Exception("Alignment of 1 sequence have no sense\n")
    protein_filename = "proteins.fasta"
    aln_filename = 'alignment.aln'

    proteins = []
    i = 1
    for seq_record in seq_records:
      protein_id = str(i) + "_" + seq_record.id
      i += 1
      proteins.append(SeqRecord(seq=Seq.translate(seq_record.seq), id=protein_id))
    with open(protein_filename, 'w') as protein_file:
        SeqIO.write(proteins, protein_file, 'fasta')

    aligner = ClustalwCommandline('clustalw', infile=protein_filename, matrix='BLOSUM', outfile=aln_filename)
    for_display, stderr = aligner()

    with open(aln_filename) as aln_file:
      with open(phylip_filename, 'w') as phylip_file:
        alignment = AlignIO.parse(aln_file, 'clustal')
        AlignIO.write(alignment, phylip_file, 'phylip')

    os.remove(protein_filename)
    return for_display
    return 0

def main(filenames):
  if len(filenames) == 0:  
    sys.stderr.write("USAGE: task2.py file1 file2 ...\n")
    return -1

  print "\nReading files..."
  seq_records = read_seqs(sys.argv[1:]) 

  print "\n\nGC-content calculation..."
  for seq_record in seq_records:
    print "GC-content of " + seq_record.id + " = " + get_gc_content(seq_record) + "%"

  print "\nTranscription..."
  for seq_record in seq_records:
    print "transcription of " + seq_record.id + ":"
    print transcription(seq_record) + '\n'
    
  print "\nTranslation..."
  for seq_record in seq_records:
    print "translation of " + seq_record.id + ":"
    print translation(seq_record) + '\n'

  print "\nMultiple alignment..."
  try:
   print multiple_align(seq_records, "alignment.phy")
  except Exception as e:
    sys.stderr.write(e.message)
   
  return 0

if __name__ == "__main__":
  exit(main(sys.argv[1:]))
